from maix import ext_dev, pinmap, err, time, touchscreen, app, display, image
import math

disp = display.Display()
img = image.Image(disp.width(), disp.height())

QMI8658_I2CBUS_NUM = 4

imu = ext_dev.qmi8658.QMI8658(QMI8658_I2CBUS_NUM,
                              mode=ext_dev.imu.Mode.DUAL,
                              acc_scale=ext_dev.imu.AccScale.ACC_SCALE_2G,
                              acc_odr=ext_dev.imu.AccOdr.ACC_ODR_8000,
                              gyro_scale=ext_dev.imu.GyroScale.GYRO_SCALE_16DPS,
                              gyro_odr=ext_dev.imu.GyroOdr.GYRO_ODR_8000)

# 初始化角度和时间
roll_angle = 0
pitch_angle = 0
yaw_angle = 0
last_time = time.ticks_ms()

# 卡尔曼滤波器参数
Q_angle = 0.001  # 角度的估计误差
Q_gyro = 0.001   # 陀螺仪的估计误差
R_angle = 0.01   # 角度的观测噪声
x = 0            # 角度的估计值
P = 0.01            # 估计值的协方差

# 低通滤波器参数
alpha_lowpass = 0.4
last_acc_roll = 0
last_acc_pitch = 0

def kalman_filter(z, dt):
    global x, P
    # 预测阶段
    x = x + dt * (data[3] + data[4]) / 2  # 用陀螺仪数据预测角度变化
    P = P + Q_angle - Q_gyro * dt

    # 更新阶段
    K = P / (P + R_angle)
    x = x + K * (z - x)
    P = (1 - K) * P

    return x



while True:
    current_time = time.ticks_ms()
    data = imu.read()

    # 计算时间间隔（毫秒转换为秒）
    dt = (current_time - last_time) / 1000.0
    
    # 计算陀螺仪的角度（度每秒 * 秒）
    roll_gyro = roll_angle + data[3] * dt
    pitch_gyro = pitch_angle + data[4] * dt
    yaw_gyro = yaw_angle + data[5] * dt 

    # 计算加速度计的角度
    acc_roll = math.atan2(data[1], data[2]) * 180 / math.pi
    acc_pitch = math.atan2(data[0], math.sqrt(data[1]**2 + data[2]**2)) * 180 / math.pi

    # 应用低通滤波器
    acc_roll = alpha_lowpass * acc_roll + (1 - alpha_lowpass) * last_acc_roll
    acc_pitch = alpha_lowpass * acc_pitch + (1 - alpha_lowpass) * last_acc_pitch

    # 应用卡尔曼滤波器
    roll_angle = kalman_filter(acc_roll, dt)
    pitch_angle = kalman_filter(acc_pitch, dt)
    
    # 更新偏航角
    yaw_angle = yaw_gyro  # 更新偏航角

    # 更新低通滤波器的上一次值
    last_acc_roll = acc_roll
    last_acc_pitch = acc_pitch

    # 显示IMU信息
    img.clear()
    img.draw_string(10, 10, f"Acc X: {data[0]:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    img.draw_string(10, 40, f"Acc Y: {data[1]:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    img.draw_string(10, 70, f"Acc Z: {data[2]:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    img.draw_string(10, 100, f"Gyro X: {data[3]:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    img.draw_string(10, 130, f"Gyro Y: {data[4]:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    img.draw_string(10, 160, f"Gyro Z: {data[5]:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    img.draw_string(10, 190, f"Temp: {data[6]:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    
    # 显示计算的角度
    img.draw_string(10, 220, f"Roll Angle: {roll_angle:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    img.draw_string(10, 250, f"Pitch Angle: {pitch_angle:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    img.draw_string(10, 280, f"Yaw Angle: {yaw_angle:.2f}", image.Color.from_rgb(255, 0, 0), 2)
    
    disp.show(img)
    
    last_time = current_time
